//
// Created by wangwanlong on 2024/11/1.
//

#ifndef MKAHYPAR_ML_COARSENER_H
#define MKAHYPAR_ML_COARSENER_H

#include "mkahypar/partition/coarsening/i_coarsener.h"
#include "mkahypar/partition/coarsening/coarsener_base.h"
#include "mkahypar/partition/coarsening/vertex_pair_rater.h"
#include "mkahypar/utils/timer.h"
#include "mkahypar/partition/metrics.h"
//重边缩合
namespace mkahypar {
  template<class ScorePolicy = HeavyEdgeScore,
          class HeavyNodePenaltyPolicy = NoWeightPenalty,
          class CommunityPolicy = UseCommunityStructure,
          class RatingPartitionPolicy = NormalPartitionPolicy,
          class AcceptancePolicy = BestRatingPreferringUnmatched<>,
          typename RatingType = RatingType>
  class MLCoarsener final : public ICoarsener, CoarsenerBase {
  private:
    //打分函数
    using Rater = VertexPairRater<ScorePolicy,
            HeavyNodePenaltyPolicy,
            CommunityPolicy,
            RatingPartitionPolicy,
            AcceptancePolicy,
            RatingType>;
    using Rating = typename Rater::Rating;
    static constexpr HypernodeID kInvalidTarget = std::numeric_limits<HypernodeID>::max();
  public:
    MLCoarsener(Hypergraph &hypergraph, TargetGraph &topo) :
            CoarsenerBase(hypergraph, topo), _rater(hypergraph) {}

    //执行还原
    void uncontract() {
      restoreParallelHyperedges();
      restoreSingleNodeHyperedges();
//      LOG<<"uncontract"<<_history.back().contraction_memento.u<<_history.back().contraction_memento.v;
      _hg.uncontract(_history.back().contraction_memento);
      _history.pop_back();

    }

    //执行还原
    void uncontract(Metrics &best) {
//      LOG << "check1";
//      check();
      //还原平行边
      _hypergraph_pruner.restoreParallelHyperedges(_hg, _history.back(), best, _topo);
//      LOG << "check2";
//      check();
      restoreSingleNodeHyperedges();
//      LOG << "check3";
//      check();
//      LOG<<"uncontract"<<_history.back().contraction_memento.u<<_history.back().contraction_memento.v;
      _hg.uncontract(_history.back().contraction_memento);
      _history.pop_back();
//      LOG << "check4";
//      check();

//      _hg.printGraph();

    }

    void check() {
      //检查超边的数据结构
      int num_partitions = _hg._k;
      for (const HyperedgeID &he: _hg.edges()) {
        int edge_size = 0;
        auto source = _hg.hyperedge(he).source();
        std::vector<int> pin_in_part_size(num_partitions);
        std::vector<int> part_route_count(num_partitions);
        int64_t connect = 0;
        for (auto hn: _hg.pins(he)) {
          ++edge_size;
          pin_in_part_size[_hg.partID(hn)] += 1;
          connect |= 1LL << _hg.partID(hn);
          if (hn != source)
            part_route_count[_hg.partID(hn)] += _hg.tarRouteCount(he, hn);
        }
        //检查联通集
        if (connect != _hg.hyperedge(he).connectivity._n) {
          LOG << "connectivity  error!";
          LOG << he << connect << _hg.hyperedge(he).connectivity._n;
          ERROREXIT("");
        }
        //检查edge size
        if (edge_size != _hg.edgeSize(he)) {
          LOG << "edge size error!";
          LOG << he << edge_size << _hg.edgeSize(he);
          ERROREXIT("");
        }
        for (int id = 0; id < num_partitions; ++id) {
          //检查pin size
          if (pin_in_part_size[id] != _hg.pinCountInPart(he, id)) {
            LOG << "pin in part size error!";
            LOG << he << pin_in_part_size[id] << _hg.pinCountInPart(he, id);
            ERROREXIT("");
          }
          //检查part route count
          if (part_route_count[id] != _hg.routeCountInPart(he, id)) {
            LOG << "route count in part size error!";
            LOG << he << id << part_route_count[id] << _hg.routeCountInPart(he, id);
            ERROREXIT("");
          }
        }
      }
      //检查超点的数据结构
      for (const HypernodeID &hn: _hg.nodes()) {
        int num_incident_cuts = 0;
        for (const auto &he: _hg.hypernode(hn).incidentNets()) {
          if (_hg.connectivity(he) > 1)
            ++num_incident_cuts;
        }
        if (num_incident_cuts != _hg.hypernode(hn).num_incident_cut_hes) {
          LOG << "node cut error!";
          LOG << hn << num_incident_cuts << _hg.hypernode(hn).num_incident_cut_hes;
          ERROREXIT("error!");
        }
      }

      //检查互联量
      std::vector<int> line(num_partitions);
      for (const HyperedgeID &he: _hg.edges()) {
        auto he_weight = _hg.hyperedge(he).weight();
        int sz = 0;
        if (_hg.connectivity(he) > 1) {
          for (int id = 0; id < num_partitions; ++id) {
            if (_hg.pinCountInPart(he, id) > 0)
              line[id] += he_weight;
          }
        }
      }
      for (int id = 0; id < num_partitions; ++id) {
        if (line[id] != _topo._fpga[id]._used_line) {
          LOG << "line error!";
          LOG << id << line[id] << _topo._fpga[id]._used_line;
          ERROREXIT("error!");
        }
      }
    }


    bool doUncoarsen(IRefiner &refiner, Metrics &best_obj) {
      //执行初始解
      //计算初始目标函数
      refiner.computeObj(best_obj);
//      LOG << "block violate:" << best_obj.block_penalty;
//      LOG << "edge violate:" << best_obj.edge_penalty;
//      LOG << "hop violate:" << best_obj.hop_penalty;

      //初始化后执行搜索
      std::vector<HypernodeID> active_nodes;
      for (auto hn: _hg.nodes()) {
        active_nodes.emplace_back(hn);
      }
      refiner.refine(active_nodes, best_obj);
//      best_obj.print();
      //还原
      while (!_history.empty()) {
        if (Timer::getInstance().TLE()) {
          //时间到了直接还原
          while (!_history.empty()) {
            uncontract(best_obj);
          }
          break;
        }

        //还原一对顶点
        active_nodes.clear();
        active_nodes.push_back(_history.back().contraction_memento.u);
        active_nodes.push_back(_history.back().contraction_memento.v);

        //还原平行边和单点边

        uncontract(best_obj);
        //执行搜索
        refiner.refine(active_nodes, best_obj);
//        LOG << "now hypernode num:" << _hg._num_hypernodes;
//        LOG << "now hyperedge num:" << _hg._num_hyperedges;
//        best_obj.print();

      }
//      LOG << "final seach:";
      active_nodes.clear();
      for (auto hn: _hg.nodes()) {
        active_nodes.emplace_back(hn);
      }
      refiner.refine(active_nodes, best_obj);
//      best_obj.print();
      return true;
    }

  private:
    //执行缩合
    void coarsenImpl(const HypernodeID limit) override {
      int pass_nr = 0;
      std::vector<HypernodeID> current_hns;
      while (_hg._num_hypernodes > limit) {
        _rater.resetMatches();
        current_hns.clear();
        const HypernodeID num_hns_before_pass = _hg._num_hypernodes;//记录缩合前的顶点数
        for (const HypernodeID &hn: _hg.nodes()) {
          current_hns.push_back(hn);
        }
        Randomize::getInstance().shuffleVector(current_hns, current_hns.size());
        for (const HypernodeID &hn: current_hns) {
          if (_hg.nodeIsEnabled(hn)) {
            const Rating rating = _rater.rate(hn);

            if (rating.target != kInvalidTarget) {
              _rater.markAsMatched(hn);
              _rater.markAsMatched(rating.target);
              // if (_hg.nodeDegree(hn) > _hg.nodeDegree(rating.target)) {

              performContraction(hn, rating.target);

//              _hg.printGraph();
              // } else {
              //   contract(rating.target, hn);
              // }
            }

            if (_hg._num_hypernodes <= limit) {
              break;
            }
          }
        }
        //此轮没有找到缩合顶点对，直接退出缩合
        if (num_hns_before_pass == _hg._num_hypernodes) {
          break;
        }
        ++pass_nr;
      }
//      LOG << "coarsen finish,total pass = " << pass_nr;
//      LOG << "current node num:" << _hg._num_hypernodes;
//      LOG << "current edge num:" << _hg._num_hyperedges;

    }

    bool uncoarsenImpl(IRefiner &refiner, Metrics &best) override {
      return doUncoarsen(refiner, best);
    }

    Rater _rater;//负责顶点对打分

  };
}
#endif //MKAHYPAR_ML_COARSENER_H
